"""
角色：资料收集者research、画图者painter、团队team
画出英国过去5年的国内生产总值的折线图
1.构建llm _researcher_agent  _painter_agent 对象
2.定义节点  researcher_node painter_node
3.定义路由  router_edge  painter_node--researcher_node 、researcher_node--》painter_node 、painter_node--》END
4.定义和构建图：StateGraph(MessagesState)
5._graph.stream流的方式运行图
"""
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from model_utils import getLLM

_llm  = getLLM()

from research import Researcher
from painter import Painter

_researcher_agent = Researcher(_llm)
_painter_agent = Painter(_llm)

def researcher_node(state):
    return {"messages":[_researcher_agent(state)]}

def painter_node(state):
    return {"messages":[_painter_agent(state)]}

from langgraph.graph import StateGraph,START,END,MessagesState,MessageGraph

def router_edge(state):
    print("team router node state",state)
    messages = state["messages"]
    last_message = messages[-1]
    if "FINAL ANSWER" in last_message.content:
        return END
    return "continue"

_builder = StateGraph(MessagesState)
_builder.add_node("researcher_node",researcher_node)
_builder.add_node("painter_node",painter_node)
_builder.add_edge(START,"researcher_node")
_builder.add_conditional_edges("researcher_node",router_edge,{"continue":"painter_node",END:END})
_builder.add_conditional_edges("painter_node",router_edge,{"continue":"researcher_node",END:END})
_graph = _builder.compile()
_graph.get_graph().draw_mermaid_png(output_file_path="team.png")


# _content  = "Fetch the UK's GDP over the past 5 years,  then draw a line graph of it.  Once you code it up, finish"
_content = "获取英国过去5年的国内生产总值，然后画一个折线图,并保存到本地文件名为uk_5gdp.png，一旦你把它编码好，就执行该代码。"
# for msg in _graph.stream({"messages":[HumanMessage(content=_content)]},{"recursion_limit": 10}):
for msg in _graph.stream({"messages":[HumanMessage(content=_content)]}):
    print(msg)